import os
import shutil
import sys
import threading
import time
from pathlib import Path

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
sys.path.append(os.path.join(current_dir, "indextts"))

import gradio as gr
from utils.webui_utils import next_page, prev_page
from indextts.infer import IndexTTS
from tools.i18n.i18n import I18nAuto

i18n = I18nAuto(language="zh_CN")
MODE = 'local'
tts = IndexTTS(model_dir="checkpoints", cfg_path="checkpoints/config.yaml")

os.makedirs("outputs/tasks", exist_ok=True)
os.makedirs("outputs/batch", exist_ok=True)
os.makedirs("prompts", exist_ok=True)

mutex = threading.Lock()


def infer(voice, text, output_path=None):
    if not output_path:
        output_path = os.path.join("outputs", f"spk_{int(time.time())}.wav")
    with mutex:
        tts.infer(voice, text, output_path)
    return output_path


def gen_single(prompt, text):
    output_path = infer(prompt, text)
    return gr.update(value=output_path, visible=True)


def update_prompt_audio():
    return gr.update(interactive=True)


# ---------- 批量顺序合成 ----------
def batch_tts(prompt_audio_path, txt_folder_path, progress=gr.Progress()):
    if not prompt_audio_path:
        return "❌ 请先上传参考音频！"
    if not txt_folder_path:
        return "❌ 请输入 txt 文件夹路径！"
    folder = Path(txt_folder_path.strip())
    if not folder.is_dir():
        return f"❌ 路径不存在：{folder}"
    txt_files = sorted(folder.rglob("*.txt"))
    if not txt_files:
        return f"❌ 文件夹 {folder} 内未找到任何 txt 文件！"

    out_dir = Path("outputs/batch")
    out_dir.mkdir(exist_ok=True)
    progress(0, desc="开始批量合成...")
    for idx, txt_path in enumerate(txt_files):
        text = txt_path.read_text(encoding="utf-8").strip()
        if not text:
            continue
        wav_path = out_dir / txt_path.relative_to(folder).with_suffix(".wav")
        wav_path.parent.mkdir(parents=True, exist_ok=True)
        infer(prompt_audio_path, text, str(wav_path))
        progress((idx + 1) / len(txt_files), desc=f"已合成 {idx + 1}/{len(txt_files)}")
    return f"✅ 全部合成完成！结果保存在 {out_dir.resolve()}"


# -------------------- UI --------------------
with gr.Blocks() as demo:
    mutex = threading.Lock()
    gr.HTML('''
    <h2><center>IndexTTS: An Industrial-Level Controllable and Efficient Zero-Shot Text-To-Speech System</h2>
    <p align="center">
    <a href='https://www.xiamibaba.com'>Xiamibaba虾米巴巴</a>
    ''')

    with gr.Tab("音频生成"):
        with gr.Row():
            os.makedirs("prompts", exist_ok=True)
            prompt_audio = gr.Audio(label="请上传参考音频", key="prompt_audio",
                                    sources=["upload", "microphone"], type="filepath")
            prompt_list = os.listdir("prompts")
            default = ''
            if prompt_list:
                default = prompt_list[0]
            input_text_single = gr.Textbox(label="请输入目标文本", key="input_text_single")
            gen_button = gr.Button("生成语音", key="gen_button", interactive=True)
            output_audio = gr.Audio(label="生成结果", visible=False, key="output_audio")

        prompt_audio.upload(update_prompt_audio, inputs=[], outputs=[gen_button])
        gen_button.click(gen_single,
                         inputs=[prompt_audio, input_text_single],
                         outputs=[output_audio])

    with gr.Tab("批量顺序合成"):
        gr.Markdown("1. 先在上传框把**参考音频**传好（与单条共用）<br>2. 在下方文本框**粘贴 txt 文件夹绝对路径**<br>3. 点击“开始批量合成”")
        with gr.Row():
            folder_path = gr.Textbox(label="txt 文件夹路径", placeholder=r"D:\ai\my_txts")
            batch_btn = gr.Button("开始批量合成", variant="primary")
        progress_tb = gr.Textbox(label="进度/结果", interactive=False, lines=2)
        batch_btn.click(batch_tts,
                        inputs=[prompt_audio, folder_path],
                        outputs=progress_tb)

if __name__ == "__main__":
    demo.queue(20).launch(server_name="127.0.0.1", inbrowser=True)